import ptan
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical

class MPOAgent(ptan.agent.BaseAgent):
    def __init__(self, model, action_selector=ptan.actions.ProbabilityActionSelector(), device="cpu",
                 apply_softmax=False, preprocessor=ptan.agent.default_states_preprocessor):
        self.model = model
        self.action_selector = action_selector
        self.device = device
        self.apply_softmax = apply_softmax
        self.preprocessor = preprocessor

    @torch.no_grad()
    def __call__(self, states, agent_states=None):

        if self.preprocessor is not None:
            states = self.preprocessor(states)
            if torch.is_tensor(states):
                states = states.to(self.device)
        actions, probs = self.model.action(states)
        return np.array(actions.cpu().numpy()), [(probs)]



class FeatureExtractor(nn.Module):
    def __init__(self, input_shape):
        super(FeatureExtractor, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU()
        )


    def get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        x = x.float() / 255.0
        return self.conv(x)
    

class ActorModel(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(ActorModel, self).__init__()

        self.feature_extractor = FeatureExtractor(input_shape)

        feature_extractor_size = self.feature_extractor.get_conv_out(input_shape)
        self.fc = nn.Sequential(
            nn.Linear(feature_extractor_size, 256),
            nn.ReLU(),
            nn.Linear(256, n_actions)
        )

    def forward(self, x):
        features = self.feature_extractor(x)
        features = features.view(x.size()[0], -1)
        return F.softmax(self.fc(features), dim=1)
    

    def action(self, x):
        with torch.no_grad():
            probs = self.forward(x)
            dist = Categorical(probs=probs)
            action = dist.sample()
            prob = dist.probs
        return action, prob
    

    def get_action_prob(self, x):
        probs = self.forward(x)
        action_distribution = Categorical(probs=probs)
        return action_distribution.probs
    

    def evaluate_action(self, x, action):
        probs = self.forward(x)
        action_distribution = Categorical(probs=probs)
        log_probs = action_distribution.log_prob(action)
        entropy = action_distribution.entropy().mean()
        return action_distribution.probs, log_probs, entropy
    

class CriticModel(nn.Module):
    def __init__(self, input_shape, action_size):
        super(CriticModel, self).__init__()

        self.feature_extractor = FeatureExtractor(input_shape)

        feature_extractor_size = self.feature_extractor.get_conv_out(input_shape)
        self.fc = nn.Sequential(
            nn.Linear(feature_extractor_size, 256),
            nn.ReLU(),
            nn.Linear(256, action_size)
        )

    def forward(self, x):
        features = self.feature_extractor(x)
        features = features.view(x.size()[0], -1)
        return self.fc(features)